###############################################################################
## Main File: R functions for MCMC simulation
## Author: Xun Pang, Barry Friedman, Andrew Martin, Kevin Quinn
## This file contains the fillowing functions   
##  1.  Probit.ChangingPoint
##  2.  Marginal.likelihood 
##  3.  ReducedMCMC2
##  4.  Marginal.Probit 
##  5.  Probit.MCMC
## Caution:
##    Don't load linear model at the same time, since functions are redifined
##    After loading this file, the following functions can be called
##    However, the change point model and the probit model can run seperately
###############################################################################
  library(mvtnorm)        # library required

  ####################################
  ## Define Functions
  ####################################

  ##\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
  ## Compute the likelihood of binary response
  ##\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\

  Likelihood <- function(y, yo, theta, x)
   {
     # cat("x ",sum(x))
     if(!is.matrix(x)){
      x <- t(as.matrix(x))
     }
     mean.y <- x%*%theta
     likelihood <- yo*pnorm(mean.y)+ (1-yo)* (1-pnorm(mean.y))
     return(prod(likelihood))
   }


  ##\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
  ## Compute the likelihood of latent continuous response
  ##\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
 
  Likelihood3 <- function(y, theta, x)
   {
     k <- length(y)
     if(!is.matrix(x)){
       x <- t(as.matrix(x))
     }
     mean.y <- x%*%theta
     likelihood <- dmvnorm(y, mean=mean.y, sigma=diag(k), log=T)
     return(exp(sum(likelihood)))
   }

  ##\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
  ## Compute the predictive prob.
  ##\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\

  prediction <- function(fil.tm1, bigP, M)                                             
  {
    if(t==1) predprob = c(1,rep(0, (M-1)))           
    else predprob <- fil.tm1%*%bigP  
     return(predprob)                            
  }

  ##\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
  ## Compute the filtering prob.
  ##\\\\\\\\\\\\\\\\\\\\\\\\\\\\\

  filter.prob <- function(y, x, Theta, pred,  M)
  {
      likelihood <- sapply(c(1:M),
             function(k){Likelihood3(y=y, x=x, Theta[k,])})
      numon <- pred*likelihood
      numon <- na.omit(numon)
      demon <- sum(numon)
      if(demon==0) pst <- rep(1/M, M)
      else pst <- numon/sum(numon)
      return(pst)
  }
  
  ##\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
  ## Function: Filtering Updating
  ## Y: the response variable
  ##    in probit models, Y is the latent response
  ## X: design matrix
  ## bigP: the m-dimensional vector of transition probs.
  ## Theta: Coefficient matrix 
  ## TimeIndex: index of periods of observations 
  ##\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\

  Filtering <- function(Y, X, TimeIndex, bigP, Theta, pred.return=FALSE, M=M)
   {
    # number of states
      N <- length(unique(TimeIndex))
    # create a matrix for the filtering probability
      filtering.matrix <- matrix(NA, nrow=N, ncol=M)
      if(pred.return==TRUE){
         pred.matrix <- matrix(NA, nrow=N, ncol=M)
         pred.matrix[1,] <- c(1, rep(NA, (M-1)))
       }
      for (t in 1:N){
       filtering.prob <- rep(NA, M)
       if(t==1) predprob = c(1,rep(0, (M-1)))               
       else predprob <-filtering.matrix[(t-1),]%*%bigP  

        if(pred.return==TRUE){
          pred.matrix[t,] <- predprob
        }
       Yn <- Y[TimeIndex==t]
       Xn <- X[TimeIndex==t,]
       if (length(which(TimeIndex==t))==1) Xn <- t(Xn)
       filtering.prob <- filter.prob(y=Yn, x=Xn, Theta,  pred=predprob, M=M)
       filtering.matrix[t,] <- filtering.prob
     }
       if(pred.return==TRUE) return(list(filtering.matrix, pred.matrix))
       else return(filtering.matrix)
    }

  ##\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
  ## Function: Backward Smoothing sampling
  ##\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\

   backward.smoothing <- function(filter, bigP, N, M)
  {
     sample.s <- rep(NA,N)
     sample.s[N] <- M
     t <- N-1
     while (t >= 1){
       statet1 <- sample.s[t+1]
        if (statet1==1) sample.s[t]<-1
       else{
         pool <- c( statet1-1, statet1)
         numon <- filter[t,]*bigP[,statet1]
          if (sum(numon)==0) {
           u <- runif(1)
           probmass=c(u, (1-u))
           sample.s[t] <- sample(pool, 1, prob=probmass)
         }else{
           mass   <-  exp(log(numon)-log(sum(numon)))
           probmass <- mass[c((statet1-1),statet1)]
           sample.s[t] <- sample(pool, 1, prob=probmass)
         }
       }
       t <- t-1
     }
       return(sample.s)
   }


  ##\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
  ## Function:Frequency of staying in a state
  ##\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\

  switchg <- function (s1, M) {  # function to compute switch
    stay <- table(s1) - 1
   # leng.stay <- length(stay)
   # if(leng.stay<M) stay2 <- c(rep(0, (M-leng.stay)), stay)
   # else stay2 <- stay
   # return(stay2)'
    return(stay)
  }

  ##\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
  ## Function: Compute transition prob. matrix
  ##\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\

  trans <- function (M, PP) {
    PPP <- c(PP,1)
    trans <- diag(PPP, M)
    PPPP <- rep(1, M)-PPP
    for (i in 1:(M-1)){
     trans[i, (i+1)] <- PPPP[i]
    }
    return(trans)
  }

  ##\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
  ## Function: Data Augmentation of binary response (A-C)
  ##\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\

   BTnormlog <- function(x, mu, sigma=1){
      if(x==0) {
         u <- runif(1)
         xstar <- qnorm(log(u)+pnorm(0, mu, sigma, log=T), mu, sigma, log=T)
      }else{
         u <- runif(1)
         xstar <- -qnorm(log(u)+pnorm(0, -mu, sigma, log=T), -mu, sigma, log=T)
      }
      xstar
   }

  ##\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
  ## Function: Outer product matrix
  ##\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\

   outerself<-function(x) outer(x,x)

   ################################################################################
   ## MCMC Simulation
   ################################################################################
   Probit.ChangingPoint<-function (y,X,TimeIndex, B0, beta0, M, a00, b00,
                                  beta.start, S.start,P.start, m=10000,burnin=5000,
                                  marginal.likelihood=FALSE, tracking=50, thin=1,
                                  reduced.n=1000, reduced.burnin=500)
   {
    #check legitimacy of the starting values
    N <- length(y)
    T <- length(unique(TimeIndex))
    state.number <- length(unique(S.start))
    changing.point <- length(P.start)
    if (state.number!=(changing.point+1)){
       stop("Priors on the transition probability does not match the the prior on states. \n")  
    }
    covariate.size <- ncol(X)
    coefficient.length <- length(beta0)*M
    
    # Storage
    total.return <- m+burnin
    beta<-matrix(NA, nrow=total.return, ncol=coefficient.length)
    State <- matrix(NA, nrow=total.return, ncol=T)
    ystarm<-matrix(NA, nrow=total.return, ncol=N)
    Pvector <- matrix(NA, nrow=total.return, ncol=changing.point) 
    
    # initial values
    beta[1,] <-matrix(beta.start, nrow=1, ncol=coefficient.length)
    State[1,] <- matrix(S.start, nrow=1, ncol=T)
    ystarm[1,]<-matrix(0, nrow=1, ncol=N)
    Pvector[1,] <- matrix(P.start, nrow=1, ncol=changing.point)

    # inverting matrix
    pre.B0 <- solve(B0)
    beta.B <- pre.B0%*%beta0

    #loop begins
    mcmc <- (m+burnin)*thin
    for(g in 2:mcmc) { 
    if (tracking>0){
      if ((g%% tracking ==0)| g==(m*thin)){
         cat("g=",g,"\n")
      }
    }
    ########################################
    ## Blog 1, 2: Data Augmentation and Beta
    ########################################
    if (g==2){
      state.current <- State[1,]
      beta.current <- matrix(beta[1,], nrow=M, byrow=TRUE)
    }else{
      state.current <- State.new
      beta.current <- Beta
    }

     ystar <- rep(NA, N)
     entry <- 1:N
     Beta <- matrix(NA, ncol=covariate.size, nrow=M)
     for (i in 1:M){
         group <- which(state.current==i)
         TT <- length(group)
         if (TT==0) betadrawn <- rmvnorm(1, beta0, B0)
         else{
           beta.state <- beta.current[i,]  
           XX <- data.frame(X)[TimeIndex%in%group,]
           YY <- y[TimeIndex%in%group]
           indexY <- entry[TimeIndex%in%group]
           NN <- length(YY)
           total<-0
           ystar.M <- rep(NA, NN)
           for (k in 1:NN) {
             mu <- as.matrix(XX[k,]) %*% cbind(beta.state) 
             ystar.M[k] <- BTnormlog(x=YY[k], mu=mu)
             total<-total+outerself(as.matrix(XX[k,])[1,])
           }

           B1<-solve(total+pre.B0)
           beta1<-B1%*%(t(XX)%*%ystar.M+beta.B)
           betadrawn<-rmvnorm(1, beta1, B1)
           ystar[indexY] <- ystar.M
         }
          Beta[i,] <- betadrawn
    }
 
    ######################################
    ## Block 3: transition prob. updating
    ######################################
      
    stay.state <-  switchg(state.current, M=M)  
    bigP.draw <- c()
    for (i in 1:(M-1)){
      bigP.draw[i] <- rbeta(1, (a00+stay.state[i]), (b00+1))
    }
    bigP <-  trans (M, PP=bigP.draw)  
    
    ######################################
    ## Block 4: Update the states
    ######################################
    filtering.prob <- Filtering(Y=ystar, X=X, TimeIndex, bigP=bigP, Theta=Beta, 
                                pred.return=FALSE, M=M)  
    State.new <- backward.smoothing(filter=filtering.prob, bigP=bigP,
                                    N=T, M=M)

    ###############################################
    # store the MCMC draws of the current iteration
    ###############################################
     if ((g%%thin ==0)| g==mcmc){
        gg <- g%/%thin
        ystarm[gg,] <- ystar
        beta[gg,] <-  as.vector(t(Beta))
        Pvector[gg,] <- bigP.draw
        State[gg,] <- State.new
     }
  }
    ####################################
    ## Return MCMC output
    ####################################

    MCMCoutput <- list(beta=beta[-c(1:burnin),],augmentedData=ystarm[-c(1:burnin),],
                       TransitionProb=Pvector[-c(1:burnin),], states=State[-c(1:burnin),])
    if (marginal.likelihood==TRUE){
        source("ChgPtProbitMarginalLikelihood.R")
          MarLik <- Marginal.likelihood(MCMCoutput=MCMCoutput, TimeIndex=TimeIndex, beta0=beta0,
                                       B0=B0, M=M, a00=a00, b00=b00, iteration=reduced.n,
                                       burninstage=reduced.burnin, y=y, X=X)
          return(list(MCMCoutput=MCMCoutput, MarginalLikelihood=MarLik))
     }else{
          return(MCMCoutput)
     }
  }

                   

y=voting;  X=X; m=10000; B0=diag(20, 25);
                                          burnin=1000; beta0=rep(0, 25); beta.start=0;S.start=SStar;
                                          M=2; P.start=0.5;a00=1; b00=1; reduced.n=500; reduced.burnin=100;
                                           TimeIndex=TimeIndex; tracking=10; thin=1
